{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-03-02T08:39:53.894847Z",
     "start_time": "2020-03-02T08:39:53.889691Z"
    }
   },
   "outputs": [],
   "source": [
    "import argparse\n",
    "import math\n",
    "import time\n",
    "import pickle\n",
    "import torch\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.nn.parameter import Parameter\n",
    "from torch.nn.functional import gumbel_softmax\n",
    "from utils import gumbel_softmax_3d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-03-02T08:39:54.397538Z",
     "start_time": "2020-03-02T08:39:54.386294Z"
    }
   },
   "outputs": [],
   "source": [
    "# 梯度优化算法类\n",
    "class my_parameters(nn.Module):\n",
    "    \n",
    "    def __init__(self, batch_size=128, n=1024, device='cuda:0'):    \n",
    "        super(my_parameters, self).__init__()\n",
    "        self.batch_size = batch_size\n",
    "        self.n = n\n",
    "        rand = torch.randn(batch_size, n, 1, device=device) * 1e-5\n",
    "        self.init_value = rand\n",
    "        self.pps = Parameter(self.init_value)\n",
    "        \n",
    "    def expand_tensor(self, x):\n",
    "        return x.unsqueeze(0).repeat(self.batch_size, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-03-02T08:48:44.929400Z",
     "start_time": "2020-03-02T08:48:44.917692Z"
    }
   },
   "outputs": [],
   "source": [
    "class SKModel():\n",
    "    def __init__(self, n=20, seed=0, device='cuda:0'):\n",
    "        self.n = n\n",
    "        self.seed = seed\n",
    "        self.device = device\n",
    "        if self.seed > 0:\n",
    "            torch.manual_seed(self.seed)\n",
    "        self.J = torch.randn(self.n, self.n) / math.sqrt(self.n)\n",
    "        self.J = torch.triu(self.J, diagonal=1)\n",
    "        self.J = self.J + self.J.t()\n",
    "        self.J_np = self.J.numpy()\n",
    "\n",
    "    def gumbel_optimization(self, bs=128, max_iters=20000, lr=1, eta=1e-3, init_tau=20, final_tau=1, T=100,u=8):\n",
    "        n = self.n\n",
    "        device = self.device\n",
    "        J = self.J.to(device)\n",
    "        best_log = []\n",
    "        # learnable parameters\n",
    "        # x = torch.randn(bs, n, 1, device=device, requires_grad=True)\n",
    "        nnn = my_parameters(bs, n, device=device)\n",
    "        #x = torch.randn(bs, n, 1, device=device) * 1e-5\n",
    "        #x.requires_grad = True\n",
    "        optimizer = optim.Adam(nnn.parameters(), lr=lr)\n",
    "        tau = init_tau\n",
    "        diff=1e-8\n",
    "        decay = (init_tau - final_tau) / max_iters\n",
    "        E_best = torch.ones(1, device=device)\n",
    "\n",
    "        for i in range(max_iters):\n",
    "            E_old = E_best.clone()\n",
    "            optimizer.zero_grad()\n",
    "            probs = torch.empty(bs, n, 2, device=device)\n",
    "            p = torch.sigmoid(nnn.pps)\n",
    "            probs[:, :, 0] = p.squeeze()\n",
    "            probs[:, :, -1] = 1 - probs[:, :, 0]\n",
    "            logits = torch.log(probs + 1e-10)\n",
    "            s = 2 * gumbel_softmax_3d(logits, tau=tau, hard=False)[:, :, 0] - 1\n",
    "            E = -0.5 * torch.sum((s @ J) * s, dim=1) / n\n",
    "            constraint = torch.sum(nnn.pps ** 2) / n / bs\n",
    "            loss = torch.mean(E) + eta * constraint\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            tau -= decay\n",
    "            with torch.no_grad():\n",
    "                if i % 100 == 0:                    \n",
    "                    print(i,'Current best result: %.8f' % (E_best.cpu().numpy()))\n",
    "                s = 2 * gumbel_softmax_3d(logits, tau=tau, hard=True)[:, :, 0] - 1\n",
    "                E = -0.5 * torch.sum((s @ J) * s, dim=1) / n\n",
    "                Emin = torch.min(E)\n",
    "                Emean = torch.mean(E)\n",
    "                if Emin < E_best:\n",
    "                    E_best = Emin\n",
    "                \n",
    "                ##每个周期都记录最好的10个个体\n",
    "                #curr_min_loss, mindx = torch.min(E, 0)\n",
    "                #if len(best_log) < 10:\n",
    "                #    output = nnn.pps[mindx:(mindx+1),:].data.cpu().numpy()\n",
    "                #    best_log.append([output, curr_min_loss.item()])\n",
    "                #else:\n",
    "                #    all_fitness = sorted([[vvv[1],i] for i,vvv in enumerate(best_log)])\n",
    "                #    if curr_min_loss < all_fitness[-1][0]:\n",
    "                #        #print(curr_min_loss)\n",
    "                #        del best_log[all_fitness[-1][1]]\n",
    "                #        output = nnn.pps[mindx:(mindx+1),:].data.cpu().numpy()\n",
    "                #        best_log.append([output, curr_min_loss.item()])\n",
    "                        \n",
    "                # 演化算法\n",
    "                if i % T == 0 and i > 0:\n",
    "                    # 找到最好个体的下标\n",
    "                    mindx = torch.argsort(E, dim=0,descending=False)\n",
    "                    # 将种群按照适应度从小到大排序\n",
    "                    maxdx = torch.argsort(E, dim=0,descending=True)\n",
    "                    for j in range(bs//u):\n",
    "                        #对1/8个体进行循环\n",
    "                        temp=nnn.pps.data[mindx[j], :, 0]\n",
    "                        if np.random.randn() < -3:\n",
    "                            temp = torch.randn(n) * 1e-5\n",
    "                        #找出其中一个最差的个体，用temp替换掉。temp可以是当前个体中最好的，也可以是一个随机个体\n",
    "                        nnn.pps.data[maxdx[j], :, 0]=temp\n",
    "                        #print(nnn.pps.data.size())\n",
    "                        #print(temp.size())\n",
    "\n",
    "                  \n",
    "                \n",
    "                if torch.abs(Emin - E_old) < diff:\n",
    "                    break\n",
    "        return E_best.cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-03-02T08:54:16.161997Z",
     "start_time": "2020-03-02T08:53:37.115623Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 Current best result: 1.00000000\n",
      "100 Current best result: -0.10321082\n",
      "200 Current best result: -0.10842466\n",
      "300 Current best result: -0.11693902\n",
      "400 Current best result: -0.15015912\n",
      "500 Current best result: -0.31190103\n",
      "600 Current best result: -0.67121553\n",
      "700 Current best result: -0.70160383\n",
      "800 Current best result: -0.72137815\n",
      "900 Current best result: -0.73664403\n",
      "1000 Current best result: -0.74092221\n",
      "1100 Current best result: -0.74374402\n",
      "1200 Current best result: -0.74605000\n",
      "# 0 \t energy: -0.74662\n",
      "0 Current best result: 1.00000000\n",
      "100 Current best result: -0.10263160\n",
      "200 Current best result: -0.10684309\n",
      "300 Current best result: -0.11983079\n",
      "400 Current best result: -0.15728623\n",
      "500 Current best result: -0.39817771\n",
      "600 Current best result: -0.65877157\n",
      "700 Current best result: -0.69947886\n",
      "800 Current best result: -0.72242737\n",
      "900 Current best result: -0.73811698\n",
      "1000 Current best result: -0.74283576\n",
      "1100 Current best result: -0.74795538\n",
      "1200 Current best result: -0.75009668\n",
      "1300 Current best result: -0.75255030\n",
      "# 1 \t energy: -0.75303\n",
      "0 Current best result: 1.00000000\n",
      "100 Current best result: -0.09859671\n",
      "200 Current best result: -0.12730648\n",
      "300 Current best result: -0.12730648\n",
      "400 Current best result: -0.14471877\n",
      "500 Current best result: -0.31271806\n",
      "600 Current best result: -0.66450715\n",
      "700 Current best result: -0.70811641\n",
      "800 Current best result: -0.72841465\n",
      "900 Current best result: -0.74195248\n",
      "1000 Current best result: -0.74788976\n",
      "1100 Current best result: -0.75076711\n",
      "1200 Current best result: -0.75310588\n",
      "# 2 \t energy: -0.75338\n",
      "0 Current best result: 1.00000000\n",
      "100 Current best result: -0.12032874\n",
      "200 Current best result: -0.12032874\n",
      "300 Current best result: -0.12315165\n",
      "400 Current best result: -0.14838970\n",
      "500 Current best result: -0.29564238\n",
      "600 Current best result: -0.67200243\n",
      "700 Current best result: -0.70553052\n",
      "800 Current best result: -0.72107774\n",
      "900 Current best result: -0.73082209\n",
      "1000 Current best result: -0.74004656\n",
      "1100 Current best result: -0.74223578\n",
      "1200 Current best result: -0.74486589\n",
      "# 3 \t energy: -0.74607\n",
      "0 Current best result: 1.00000000\n",
      "100 Current best result: -0.10347448\n",
      "200 Current best result: -0.11791703\n",
      "300 Current best result: -0.11791703\n",
      "400 Current best result: -0.14813900\n",
      "500 Current best result: -0.28105515\n",
      "600 Current best result: -0.67462617\n",
      "700 Current best result: -0.70528287\n",
      "800 Current best result: -0.72059190\n",
      "900 Current best result: -0.73523343\n",
      "1000 Current best result: -0.74376154\n",
      "1100 Current best result: -0.75366765\n",
      "1200 Current best result: -0.75414574\n",
      "# 4 \t energy: -0.75493\n",
      "0 Current best result: 1.00000000\n",
      "100 Current best result: -0.10331763\n",
      "200 Current best result: -0.11630642\n",
      "300 Current best result: -0.11972836\n",
      "400 Current best result: -0.16158241\n",
      "500 Current best result: -0.31304461\n",
      "600 Current best result: -0.66780847\n",
      "700 Current best result: -0.70443690\n",
      "800 Current best result: -0.72750604\n",
      "900 Current best result: -0.73789424\n",
      "1000 Current best result: -0.74370265\n",
      "1100 Current best result: -0.74534088\n",
      "1200 Current best result: -0.74922228\n",
      "1300 Current best result: -0.75116670\n",
      "# 5 \t energy: -0.75151\n",
      "0 Current best result: 1.00000000\n",
      "100 Current best result: -0.10822621\n",
      "200 Current best result: -0.10822621\n",
      "300 Current best result: -0.11585633\n",
      "400 Current best result: -0.14821295\n",
      "500 Current best result: -0.29368055\n",
      "600 Current best result: -0.65133214\n",
      "700 Current best result: -0.69944751\n",
      "800 Current best result: -0.71675336\n",
      "900 Current best result: -0.72643495\n",
      "1000 Current best result: -0.73540330\n",
      "1100 Current best result: -0.74144483\n",
      "1200 Current best result: -0.74536717\n",
      "# 6 \t energy: -0.74541\n",
      "0 Current best result: 1.00000000\n",
      "100 Current best result: -0.10595658\n",
      "200 Current best result: -0.10668661\n",
      "300 Current best result: -0.12491953\n",
      "400 Current best result: -0.15020663\n",
      "500 Current best result: -0.39018393\n",
      "600 Current best result: -0.67444229\n",
      "700 Current best result: -0.71681714\n",
      "800 Current best result: -0.72789812\n",
      "900 Current best result: -0.73725486\n",
      "1000 Current best result: -0.74732804\n",
      "1100 Current best result: -0.75044161\n",
      "1200 Current best result: -0.75301385\n",
      "# 7 \t energy: -0.75399\n",
      "0 Current best result: 1.00000000\n",
      "100 Current best result: -0.10418534\n",
      "200 Current best result: -0.11994752\n",
      "300 Current best result: -0.12181939\n",
      "400 Current best result: -0.15584788\n",
      "500 Current best result: -0.27251142\n",
      "600 Current best result: -0.65721685\n",
      "700 Current best result: -0.70487487\n",
      "800 Current best result: -0.72228938\n",
      "900 Current best result: -0.73514098\n",
      "1000 Current best result: -0.74180818\n",
      "1100 Current best result: -0.74579847\n",
      "1200 Current best result: -0.74867749\n",
      "1300 Current best result: -0.75181198\n",
      "# 8 \t energy: -0.75190\n",
      "0 Current best result: 1.00000000\n",
      "100 Current best result: -0.09575982\n",
      "200 Current best result: -0.12324820\n",
      "300 Current best result: -0.12398985\n",
      "400 Current best result: -0.14800307\n",
      "500 Current best result: -0.34989873\n",
      "600 Current best result: -0.67898989\n",
      "700 Current best result: -0.71627098\n",
      "800 Current best result: -0.73015022\n",
      "900 Current best result: -0.74010849\n",
      "1000 Current best result: -0.74735951\n",
      "1100 Current best result: -0.75165272\n",
      "1200 Current best result: -0.75606328\n",
      "# 9 \t energy: -0.75630\n",
      "Batch Gumbel-softmax N: 1024, bs: 128, instances: 10, mean: -0.75132, std: 0.00390, sem: 0.00123\n",
      "Running time: 39.03412 s \n",
      "\n"
     ]
    }
   ],
   "source": [
    "def main():\n",
    "    # settings\n",
    "    parser = argparse.ArgumentParser('Gumbel-softmax optimizing SK model energy')\n",
    "    parser.add_argument('--n', type=int, default=1024,\n",
    "                        help='size (default: 1024)')\n",
    "    parser.add_argument('--bs', type=int, default=128,\n",
    "                        help='batch size (default: 128)')\n",
    "    parser.add_argument('--max_iters', type=int, default=2000,\n",
    "                        help='iterations (default: 50000)')\n",
    "    parser.add_argument('--lr', type=float, default=1.,\n",
    "                        help='learning rate (default: 1)')\n",
    "    parser.add_argument('--eta', type=float, default=1e-3,\n",
    "                        help='weight decay (default: 1e-3)')\n",
    "    parser.add_argument('--init-tau', type=float, default=20.,\n",
    "                        help='initial tau in Gumbel-softmax (default: 20)')\n",
    "    parser.add_argument('--final-tau', type=float, default=1.,\n",
    "                        help='final tau in Gumbel-softmax (default: 1)')\n",
    "    parser.add_argument('--instances', type=int, default=10,\n",
    "                        help='number of ensembles (default: 1024)')\n",
    "    parser.add_argument('--device', type=str, default='cuda:0',\n",
    "                        help='cuda device (default: cuda:0)')\n",
    "    parser.add_argument('--T', type=int, default=100,\n",
    "                        help='evolution cycle (default: 100)')\n",
    "    parser.add_argument('--u', type=int, default=8,\n",
    "                        help='proportion (default: 8)')\n",
    "    args = parser.parse_args(args=[])\n",
    "    n = args.n  # size\n",
    "    bs = args.bs # batch size\n",
    "    max_iters = args.max_iters\n",
    "    init_tau = args.init_tau\n",
    "    final_tau = args.final_tau\n",
    "    lr = args.lr\n",
    "    eta = args.eta\n",
    "    instances = args.instances\n",
    "    device = args.device\n",
    "    T = args.T\n",
    "    u = args.u\n",
    "    torch.manual_seed(2050)\n",
    "\n",
    "    # training\n",
    "    results_arr = []\n",
    "    for _ in range(instances):\n",
    "        sk = SKModel(n, device=device)\n",
    "        energy = sk.gumbel_optimization(bs, max_iters, lr, eta, init_tau, final_tau,T,u)\n",
    "        print('# %i \\t energy: %.5f' % (_, energy))\n",
    "        results_arr.append(energy)\n",
    "\n",
    "   # print mean energy and std\n",
    "    data = np.array(results_arr)\n",
    "    data_mean = data.mean()\n",
    "    data_std = np.sqrt(np.var(data, ddof=1))\n",
    "    data_sem = data_std / np.sqrt(instances)\n",
    "    print('Batch Gumbel-softmax N: %i, bs: %i, instances: %i, mean: %.5f, std: %.5f, sem: %.7f' %\n",
    "          (n, bs, instances, data_mean, data_std, data_sem))\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    start = time.time()\n",
    "    main()\n",
    "    end = time.time()\n",
    "    print('Running time: %.5f s \\n' % (end - start))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "oldHeight": 122.4,
   "position": {
    "height": "144px",
    "left": "397.2px",
    "right": "20px",
    "top": "121px",
    "width": "350px"
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "varInspector_section_display": "block",
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
